package com.newway.config;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    StringRedisTemplate stringRedisTemplate;

    public static final String DEFAULT_KEY_PREFIX = "spring_ai_alibaba_chat_memory:";



    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = DEFAULT_KEY_PREFIX + conversationId;
        String result = JSON.toJSONString(messages);
        System.out.println("add------------------------------:"+result);
        stringRedisTemplate.opsForList().rightPush(key,result);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key = DEFAULT_KEY_PREFIX + conversationId;
        List<String> messageStrings = stringRedisTemplate.opsForList().range(key, -lastN, -1);
        System.out.println("get-----------------------------------------------------:"+messageStrings.toString());
        List<Message> result = new ArrayList<>();
        for(String str: messageStrings){
            JSONArray objects = JSON.parseArray(str);
            JSONObject temp = (JSONObject)objects.get(0);
            String messageType = temp.getString("messageType");
            Message message = null;
            if(MessageType.USER.toString().equals(messageType)){
                message = new UserMessage(temp.getString("content"));
            }else if(MessageType.ASSISTANT.toString().equals(messageType)){
                message = new AssistantMessage(temp.getString("content"));
            }
            if (message!=null){
                result.add(message);
            }
        }
        return result;
    }

    @Override
    public void clear(String conversationId) {
        String key = DEFAULT_KEY_PREFIX + conversationId;
        stringRedisTemplate.delete(key);
    }
}
